import os
import sys

from openai import OpenAI
sys.path.append('/JiuTian/swarm/swarm-main/swarm')
from swarm import Swarm,Agent

#测试基本功能
client = OpenAI(api_key=os.environ["OPENAI_API_KEY_JIUTIAN"], base_url=os.environ["OPENAI_BASE_URL_JIUTIAN"])

def transfer_to_agent_b():
    return agent_b
#测试多智能体功能
swarm_client = Swarm(client)
agent_a = Agent(name = "Agent A",instructions= "你是一个乐于助人的助手",functions=[transfer_to_agent_b])
agent_b = Agent(name = "Agent B",instructions= "只使用儿歌回答。")

response = swarm_client.run(agent=agent_a,messages=[{"role":"user","content":"我想和agent_b谈话"}])
print(response.messages[-1]["content"])